#include <string.h>
#include <stdlib.h>

#include "tdp2.hpp"
#include "tdp2_rc.hpp"
#include "tdp2_rc5.hpp"
#include "tdp2_rc5_sm.hpp"

using namespace TDP;
using namespace TDP::Compression;
using namespace TDP::Compression::RC5;
using namespace TDP::Compression::RC5::StructuredModel;

#define FLUSH_RLE	\
	do \
	{ \
		if(inRun) \
		{ \
			runLength--; \
			valuesOut += runLength; \
			while(runLength)\
			{ \
				*outc = ReadMTF(0, indexes, last); \
				last = 0; \
				outc++; \
				runLength--; \
			} \
			inRun = 0; \
		} \
	} while(0)


inline UInt8 ReadMTF(UInt8 original, UInt8 *indexes, UInt8 last)
{
	UInt8 result;
	UInt8 temp;

	result = indexes[original];

	if((MTFBehavior_k == MTFB_NoSortOn1_k || MTFBehavior_k == MTFB_NoSortOn1LastNonZero_k) && original == 1)
	{
		if(MTFBehavior_k == MTFB_NoSortOn1_k || last != 0)
		{
			temp = indexes[0];
			indexes[0] = indexes[1];
			indexes[1] = temp;
		}
	}
	else if(original)
	{
		for(int i=original;i>=1;i--)
			indexes[i] = indexes[i-1];
		indexes[1] = result;
	}

	return result;
}

static UInt32 FillRC(Compression::RangeCoder::Decoder *dec, const void *input, UInt32 &bytesRemaining)
{
	UInt32 nBytes = static_cast<UInt32>(dec->GetFillCount());
	const UInt8 *b = reinterpret_cast<const UInt8 *>(input);
	if(nBytes > bytesRemaining)
	{
		UInt8 fillBuf[5] = { 0, 0, 0, 0, 0 };

		for(UInt32 i=0;i<bytesRemaining;i++)
			fillBuf[i] = b[i];
		dec->Fill(fillBuf);
		bytesRemaining = 0;
		return bytesRemaining;
	}
	else
	{
		bytesRemaining -= nBytes;
		dec->Fill(b);
	}
	return nBytes;
}

inline UInt32 ReadRC(RangeCoder::Decoder *rc, RangeCoder::SimpleProbabilityTable *ptable, UInt32 *outsym, const UInt8 *input, UInt32 &bytesRemaining)
{
	UInt32 ppoint = rc->GetThreshold(ptable->probTotal);

	UInt32 low = 0;
	UInt32 high = 0;
	UInt32 range;
	UInt32 i;
	for(i=0;;i++)
	{
		low = high;
		range = ptable->probabilities[i];
		high += range;

		if(ppoint >= low && ppoint < high)
			break;
	}
	*outsym = i;

	rc->Decode(low, range);
	return FillRC(rc, input, bytesRemaining);
}


inline UInt32 ReadSM(RangeCoder::Decoder *rc, StructuredModel::ProbabilityPool *mset, UInt16 *outsym, const UInt8 *input, UInt32 &bytesRemaining)
{
	const UInt8 *inBase;
	UInt32 secondModel = 0;
	UInt32 sym;
	UInt32 addon;

	RangeCoder::SimpleProbabilityTable ptable;

	inBase = input;

	ptable.probabilities = mset->probPool;
	ptable.probTotal = mset->ptotals[0];
	input += ReadRC(rc, &ptable, &sym, input, bytesRemaining);
	mset->UpdateProbabilities(0, sym);

	if(sym == ValRun0_k)		{ *outsym = 0; return input - inBase; }
	else if(sym == ValRun1_k)	{ *outsym = 1; return input - inBase; }
	else if(sym == ValRun1_k)	{ *outsym = 2; return input - inBase; }
	else
	{
		if(sym == Val128_255_k)			{ addon = 129;	secondModel = M_128_255_k; }
		else if(sym == Val64_127_k)		{ addon = 65;	secondModel = M_64_127_k; }
		else if(sym == Val32_63_k)		{ addon = 33;	secondModel = M_32_63_k; }
		else if(sym == Val16_31_k)		{ addon = 17;	secondModel = M_16_31_k; }
		else if(sym == Val8_15_k)		{ addon = 9;	secondModel = M_8_15_k; }
		else if(sym == Val4_7_k)		{ addon = 5;	secondModel = M_4_7_k; }
		else /*if(sym == Val2_3_k)*/	{ addon = 3;	secondModel = M_2_3_k; }
	}

	if(secondModel)
	{
		ptable.probabilities = mset->probPool + ProbPoolOffsets_k[secondModel];
		ptable.probTotal = mset->ptotals[secondModel];
		input += ReadRC(rc, &ptable, &sym, input, bytesRemaining);
		*outsym = sym + addon;
		mset->UpdateProbabilities(secondModel, sym);
	}

	return input - inBase;
}

static bool GetBWTValues(const void *input, void *output, UInt32 csize, UInt32 dsize, UInt32 rleValues)
{
	UInt32 runLength;
	UInt32 bytesRead, bytesTotal;
	UInt32 valuesOut;
	int inRun = 0;
	UInt16 sym;
	UInt8 last=255;
	UInt8 current;

	UInt8 indexes[256];

	UInt8 *outc;
	const UInt8 *inc;

	Compression::RangeCoder::Decoder rc;
	Compression::RC5::StructuredModel::ProbabilityPool set;

	bytesTotal = valuesOut = 0;

	inc = reinterpret_cast<const UInt8 *>(input);
	outc = reinterpret_cast<UInt8 *>(output);

	set.Init();
	rc.Init();

	inc += FillRC(&rc, input, csize);

	for(int i=0;i<256;i++)
		indexes[i] = i;

	while(rleValues)
	{
		// Grab a raw symbol
		bytesRead = ReadSM(&rc, &set, &sym, inc, csize);
		bytesTotal += bytesRead;
		inc += bytesRead;

		if(bytesTotal > csize) return false;

		if(sym <= 1)
		{
			if(inRun)
				runLength = (runLength << 1) + sym;
			else
			{
				inRun = 1;
				runLength = 2 + sym;
			}

			if(dsize - valuesOut < runLength) return false;
		}
		else
		{
			FLUSH_RLE;
			if(dsize - valuesOut < runLength) return false;

			current = sym-1;
			*outc = ReadMTF(current, indexes, last);
			last = current;
			outc++;
		}

		rleValues--;
	}

	if(dsize - valuesOut < runLength) return false;
	FLUSH_RLE;

	return true;
}


static void IBWT(const void *input, void *bwtOut, UInt32 *indexes, UInt32 count, UInt32 pi)
{
	UInt32 counts[256];
	UInt32 offsets[256];

	const UInt8 *inc;
	UInt8 *outc;

	for(int i=0;i<256;i++)
		counts[i] = 0;

	// Count characters
	inc = reinterpret_cast<const UInt8 *>(input);
	for(UInt32 i=0;i<count;i++)
	{
		counts[*inc]++;
		inc++;
	}

	// Convert to offsets
	offsets[0] = 0;
	for(int i=1;i<256;i++)
		offsets[i] = offsets[i-1] + counts[i-1];

	// Create transformation vector
	inc = reinterpret_cast<const UInt8 *>(input);
	for(UInt32 i=0;i<count;i++)
	{
		indexes[offsets[*inc]++] = i;
		inc++;
	}

	// Untransform
	inc = reinterpret_cast<const UInt8 *>(input);
	outc = reinterpret_cast<UInt8 *>(bwtOut);
	for(UInt32 i=0;i<count;i++)
	{
		*outc = inc[pi];
		outc++;
		pi = indexes[pi];
	}
}



bool DecompressBlock(const void *input, void *tempspace, UInt32 csize, UInt32 dsize)
{
	const UInt8 *cinput;
	UInt32 pindex, rleValues;

	UInt8 *bwtOut;
	UInt8 *mtfrleOut;
	UInt32 *indexes;

	if(dsize > csize)
		return false;
	if(dsize == csize)
	{
		memcpy(tempspace, input, csize);
		return true;
	}

	bwtOut = reinterpret_cast<UInt8 *>(tempspace);
	mtfrleOut = bwtOut + dsize;

	if(dsize & 1)
		indexes = reinterpret_cast<UInt32 *>(mtfrleOut + dsize + 2);
	else
		indexes = reinterpret_cast<UInt32 *>(mtfrleOut + dsize);

	cinput = reinterpret_cast<const UInt8 *>(input);
	pindex = (cinput[0] << 16) | (cinput[1] << 8) | cinput[2];
	if(pindex >= dsize) return false;
	rleValues = (cinput[3] << 16) | (cinput[4] << 8) | cinput[5];
	if(rleValues > dsize || rleValues == 0) return false;

	if(!GetBWTValues(cinput+6, mtfrleOut, csize-6, dsize, rleValues))
		return 1;

	IBWT(mtfrleOut, bwtOut, indexes, dsize, pindex);

	return 1;
}

